import torch
from torchvision import models, transforms
from PIL import Image
import torch.nn.functional as F

# 加载预训练的resnet18模型
model = models.resnet18(pretrained=True)

# 删除最后一层全连接层，因为我们只关心特征，而不进行分类
model = torch.nn.Sequential(*list(model.children())[:-1])

# 将模型设置为评估模式，这是因为在推理期间我们不需要梯度计算和dropout等操作
model.eval()

# 图像预处理
transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

# 加载图像
img1 = Image.open("baseface/bigbeauty.png")
img2 = Image.open("baseface/yangyu1.png")
# img_t = transform(img)
# batch_t = torch.unsqueeze(img_t, 0)
#
# # 提取特征
# with torch.no_grad():
#     outputs = model(batch_t)
#     # 输出是一个形状为[1, 512]的张量，表示512维的特征向量
#     print(outputs.shape)
#     # 如果需要将输出转换为numpy数组，可以使用以下语句：
#     feature_vector = outputs.squeeze().numpy()
#     print(feature_vector.shape)  # 这将输出 (512,)
#     print(feature_vector)  # 这将输出 (512,)
img1_t = transform(img1)
img2_t = transform(img2)

# 创建批次
batch_t = torch.stack((img1_t, img2_t))

# 提取特征
with torch.no_grad():
    outputs = model(batch_t)
    features = [output.squeeze() for output in outputs]

# 计算余弦相似度
similarity = F.cosine_similarity(features[0].unsqueeze(0), features[1].unsqueeze(0))
print("Cosine similarity: ", similarity.item())
